extern crate winapi;

use std::collections::HashMap;
use std::ffi::OsStr;
use std::io::Cursor;
use std::os::windows::ffi::OsStrExt;
use std::ptr::null_mut;

use base64::encode;
use image::{DynamicImage, ImageBuffer, ImageEncoder, Rgba};
use image::codecs::png::PngEncoder;
use image::ExtendedColorType;
use winapi::shared::windef::HICON;
use winapi::um::shellapi::{SHFILEINFOW, SHGetFileInfoW, SHGFI_ICON, SHGFI_LARGEICON};
use winapi::um::wingdi::{BITMAP, BITMAPINFO, BITMAPINFOHEADER, CreateCompatibleDC, DeleteDC, DIB_RGB_COLORS, GetDIBits, GetObjectW, SelectObject};
use winapi::um::winuser::{GetIconInfo, ICONINFO};

pub struct IconCache {
    cache: HashMap<String, String>,
}

impl IconCache {
    pub fn new() -> Self {
        IconCache {
            cache: HashMap::new(),
        }
    }

    pub fn get_or_insert<F>(&mut self, file_path: &str, fetch_icon: F) -> Option<String>
    where
        F: FnOnce() -> Option<String>,
    {
        if !self.cache.contains_key(file_path) {
            if let Some(icon_data) = fetch_icon() {
                self.cache.insert(file_path.to_string(), icon_data);
            }
        }
        self.cache.get(file_path).cloned()
    }
}

pub fn get_file_icon(file_path: &str) -> Option<HICON> {
    let path: Vec<u16> = OsStr::new(file_path)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let mut file_info: SHFILEINFOW = unsafe { std::mem::zeroed() };

    let result = unsafe {
        SHGetFileInfoW(
            path.as_ptr(),
            0,
            &mut file_info,
            std::mem::size_of::<SHFILEINFOW>() as u32,
            SHGFI_ICON | SHGFI_LARGEICON,
        )
    };

    if result == 0 {
        None
    } else {
        Some(file_info.hIcon)
    }
}

pub fn hicon_to_base64(hicon: HICON) -> Option<String> {
    let mut icon_info: ICONINFO = unsafe { std::mem::zeroed() };
    let success = unsafe { GetIconInfo(hicon, &mut icon_info) };

    if success == 0 {
        eprintln!("Failed to get icon info");
        return None;
    }

    let mut bmp: BITMAP = unsafe { std::mem::zeroed() };
    unsafe { GetObjectW(icon_info.hbmColor as _, std::mem::size_of::<BITMAP>() as i32, &mut bmp as *mut _ as *mut _) };

    let width = bmp.bmWidth as u32;
    let height = bmp.bmHeight as u32;

    let mut bmi: BITMAPINFO = unsafe { std::mem::zeroed() };
    bmi.bmiHeader.biSize = std::mem::size_of::<BITMAPINFOHEADER>() as u32;
    bmi.bmiHeader.biWidth = width as i32;
    bmi.bmiHeader.biHeight = (height as i32);  // Top-down DIB
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 32;
    bmi.bmiHeader.biCompression = winapi::um::wingdi::BI_RGB;

    let size = (bmp.bmWidthBytes * bmp.bmHeight) as usize;
    let mut buffer = vec![0u8; size];

    unsafe {
        let hdc = CreateCompatibleDC(null_mut());
        if hdc.is_null() {
            eprintln!("Failed to create compatible DC");
            return None;
        }

        let old_bitmap = SelectObject(hdc, icon_info.hbmColor as _);

        if GetDIBits(hdc, icon_info.hbmColor as _, 0, height as u32, buffer.as_mut_ptr() as *mut _, &mut bmi, DIB_RGB_COLORS) == 0 {
            eprintln!("GetDIBits failed");
            DeleteDC(hdc);
            return None;
        }

        SelectObject(hdc, old_bitmap);
        DeleteDC(hdc);
    }

    // Convert BGR to RGB and flip image vertically
    let row_size = (width * 4) as usize; // Assuming 32-bit color
    let mut rgb_buffer = vec![0u8; size];
    for y in 0..height {
        let src_start = (height - y - 1) as usize * row_size;
        let dst_start = y as usize * row_size;
        for x in 0..width {
            let src_index = src_start + (x * 4) as usize;
            let dst_index = dst_start + (x * 4) as usize;
            // Convert BGR to RGB
            rgb_buffer[dst_index] = buffer[src_index + 2]; // Red
            rgb_buffer[dst_index + 1] = buffer[src_index + 1]; // Green
            rgb_buffer[dst_index + 2] = buffer[src_index]; // Blue
            rgb_buffer[dst_index + 3] = buffer[src_index + 3]; // Alpha
        }
    }

    let img = ImageBuffer::<Rgba<u8>, _>::from_raw(width, height, rgb_buffer)?;
    let dynamic_img = DynamicImage::ImageRgba8(img);

    let mut bytes: Vec<u8> = Vec::new();
    let mut cursor = Cursor::new(&mut bytes);

    let encoder = PngEncoder::new(&mut cursor);

    let image_buffer = dynamic_img.to_rgba8();
    if let Err(_) = encoder.write_image(&image_buffer, width, height, ExtendedColorType::Rgba8) {
        eprintln!("Failed to encode image as PNG");
        return None;
    }

    Some(encode(&bytes))
}
